#include "WorkThreadPool.h"
#include <algorithm>

namespace toolkit {

/************************ WorkThread 实现 ************************/

WorkThread::WorkThread() {
}

WorkThread::~WorkThread() {
    stop();
}

void WorkThread::start() {
    if (_thread.joinable()) {
        return;
    }
    
    _exit_flag = false;
    _thread = std::thread([this]() {
        run();
    });
    _thread_id = _thread.get_id();
}

void WorkThread::stop() {
    if (!_thread.joinable()) {
        return;
    }
    
    _exit_flag = true;
    _task_queue.push_exit(1);
    
    if (_thread.joinable()) {
        _thread.join();
    }
}

TaskCancelable::Ptr WorkThread::async(Task task, bool may_sync) {
    if (may_sync && isCurrentThread()) {
        // 当前线程是工作线程，直接执行
        task();
        return nullptr;
    }
    
    if (_exit_flag) {
        // 线程已停止
        return nullptr;
    }
    
    auto cancelable = std::make_shared<TaskCancelableImp>();
    auto wrapped_task = [cancelable, task]() {
        if (!cancelable->isCanceled()) {
            task();
        }
    };
    
    _task_queue.push_task(std::move(wrapped_task));
    return cancelable;
}

TaskCancelable::Ptr WorkThread::async_first(Task task, bool may_sync) {
    if (may_sync && isCurrentThread()) {
        // 当前线程是工作线程，直接执行
        task();
        return nullptr;
    }
    
    if (_exit_flag) {
        // 线程已停止
        return nullptr;
    }
    
    auto cancelable = std::make_shared<TaskCancelableImp>();
    auto wrapped_task = [cancelable, task]() {
        if (!cancelable->isCanceled()) {
            task();
        }
    };
    
    _task_queue.push_task_first(std::move(wrapped_task));
    return cancelable;
}

int WorkThread::getLoad() {
    return _thread_load.getLoad();
}

std::thread::id WorkThread::getThreadId() const {
    return _thread_id;
}

bool WorkThread::isCurrentThread() const {
    return std::this_thread::get_id() == _thread_id;
}

void WorkThread::run() {
    Task task;
    while (!_exit_flag) {
        // 获取任务
        if (!_task_queue.get_task(task)) {
            // 队列为空但信号量被触发，可能是退出信号
            continue;
        }
        
        // 执行任务
        _thread_load.addLoad();
        try {
            task();
        } catch (std::exception &e) {
            // 处理任务执行异常
        } catch (...) {
            // 处理未知异常
        }
        _thread_load.reduceLoad();
    }
}

/************************ WorkThreadPool 实现 ************************/

WorkThreadPool::WorkThreadPool(size_t thread_num) : _thread_num(thread_num) {
    // 创建工作线程
    for (size_t i = 0; i < thread_num; ++i) {
        auto thread = std::make_shared<WorkThread>();
        thread->start();
        _threads.emplace_back(thread);
        _thread_map[thread->getThreadId()] = thread;
    }
}

WorkThreadPool::~WorkThreadPool() {
    // 停止所有工作线程
    for (auto& thread : _threads) {
        thread->stop();
    }
    
    _threads.clear();
    _thread_map.clear();
}

size_t WorkThreadPool::size() const {
    return _thread_num;
}

TaskCancelable::Ptr WorkThreadPool::async(Task task, bool may_sync) {
    // 尝试获取当前线程对应的工作线程
    auto work_thread = getCurrentWorkThread();
    if (work_thread) {
        // 当前线程是工作线程池中的线程，直接在当前线程执行
        return work_thread->async(task, may_sync);
    }
    
    // 获取负载最小的工作线程
    work_thread = getMinLoadWorkThread();
    return work_thread->async(task, false);
}

TaskCancelable::Ptr WorkThreadPool::async_first(Task task, bool may_sync) {
    // 尝试获取当前线程对应的工作线程
    auto work_thread = getCurrentWorkThread();
    if (work_thread) {
        // 当前线程是工作线程池中的线程，直接在当前线程执行
        return work_thread->async_first(task, may_sync);
    }
    
    // 获取负载最小的工作线程
    work_thread = getMinLoadWorkThread();
    return work_thread->async_first(task, false);
}

int WorkThreadPool::getLoad() {
    int total_load = 0;
    for (auto& thread : _threads) {
        total_load += thread->getLoad();
    }
    return total_load;
}

WorkThread::Ptr WorkThreadPool::getWorkThread(size_t index) {
    if (index >= _threads.size()) {
        return nullptr;
    }
    return _threads[index];
}

WorkThread::Ptr WorkThreadPool::getCurrentWorkThread() {
    auto thread_id = std::this_thread::get_id();
    std::lock_guard<std::mutex> lock(_mutex);
    auto it = _thread_map.find(thread_id);
    if (it != _thread_map.end()) {
        return it->second;
    }
    return nullptr;
}

WorkThread::Ptr WorkThreadPool::getMinLoadWorkThread() {
    if (_threads.empty()) {
        return nullptr;
    }
    
    return *std::min_element(_threads.begin(), _threads.end(), 
        [](const WorkThread::Ptr& a, const WorkThread::Ptr& b) {
            return a->getLoad() < b->getLoad();
        });
}

} /* namespace toolkit */